[Frontend] Simulate the Triton route's kernel in TOGSim - #306
Conversation
| if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): | ||
| timing.emit_trace(self.workdir, self.meta) | ||
| result = timing.run_togsim(self.workdir) | ||
| logger.info("[triton-npu] %s simulated -> %s", self.kernel_name, result) |
There was a problem hiding this comment.
Replace triton-npu to TOGSim
There was a problem hiding this comment.
Done in 4005409 — [Spike]. Same point: this warning is about the functional half, so it should say so.
| result = timing.run_togsim(self.workdir) | ||
| logger.info("[triton-npu] %s simulated -> %s", self.kernel_name, result) | ||
| logger.warning( | ||
| "[triton-npu] %s: output tensors are NOT written; the functional " |
There was a problem hiding this comment.
triton-npu -> Spike
YWHyuk
left a comment
YWHyuk
left a comment
There was a problem hiding this comment.
The structural read is right, and it is the part that could have gone wrong: the trace producer already splits at the work-item seam, so a Triton kernel arrives in the shape the outliner would otherwise have to produce, and only the enumeration was missing. _underlying_block_arg also fixes a genuinely silent bug — arg_id = -1 breaks the address model with no diagnostic.
Four things I would change, in priority order. Two of them are about the same failure mode: a wrong number that looks like a measurement.
1. _build's new root filter is fail-open
roots = [op for op in block.operations
if op.operation.name == "affine.for" and _has_loop_role(op)]
if roots: ...
# otherwise: the whole body becomes one work-itemBefore, any top-level affine.for was a root (an untagged one fell through to _handle_compute). Now an untagged top-level loop is skipped, and if every top-level loop is untagged the kernel takes the no-loop path and yields a wrong graph — where it previously yielded an empty one, which is an obvious failure.
I checked this is latent rather than active: all eight templates (bmm / cat / conv×3 / gemm / maxpool / sdpa / sort) attach a role, and every top-level affine.for in the generated kernels I have locally (5/5) carries one.
Still, the regression it guards against is silent. Suggest: on the no-loop path, raise (or at minimum warn) if the block still contains an untagged top-level affine.for. A few lines, and it turns "cycle counts quietly drifted" into a message.
2. Placeholder cycles are invisible in the result
measure_tile_cycles returning None is logged loudly, but run_togsim's returned dict is indistinguishable from one backed by a real gem5 measurement. The cycle number is the entire output of this path, and a log line scrolls away.
Suggest carrying the fact into the result — a flag in the dict, or a marker in the TSV header that the reader surfaces — so a consumer cannot mistake PLACEHOLDER_CYCLE for a measurement. The comment on that constant already argues exactly this point; this extends it to the value that actually leaves the module.
3. work_item_for encodes the triton-shared ABI as a constant
pid_x = n_tensor + n_scalar + 3 # after gridX, gridY, gridZThe layout (pointers → user scalars → grid×3 → pid×3) is documented, but if it ever shifts, some other argument silently becomes the program id and the grid loop drives the wrong value. Suggest asserting the shape rather than only computing from it — e.g. check the trailing six arguments are i32 — so a change announces itself.
Same function assumes a 1-D grid. parallel_args scales with len(grid), but kernel_spec._grid only computes the x extent, so a second axis would make the two disagree. assert len(grid) == 1 pins that until y/z actually lands; the comment already notes x is all that is covered today.
5. _is_address_plumbing whitelists by type
The docstring owns the limitation, and the reasoning is sound for this backend — tile data is vector/float, address math is index/integer. The narrower worry is the op-name side: func.return and memref.cast are matched by name and everything else must be arith.*. If the tnpu lowering later emits an llvm.* or index.* op at top level, it is counted as vector compute instead of plumbing.
Given this file raises (_Fatal) elsewhere rather than guessing, an unknown top-level op might be better treated the same way.
Approach and framing both look good to me — stating up front what does not work (output tensors, matmul timing, 650 vs 251 from the missing double buffering) is the right way to land a partial capability. The 8192 B DRAM check is a nice touch: it verifies every dispatch in the enumeration actually ran, which the cycle count alone would not show.
Items 1 and 2 are the ones I would want before merge; 3 and 5 read fine as follow-ups.
be6802f to
466f517
Compare
It prints nothing. The method walks the IR and attaches TOG nodes; `bfs` and `display` do the printing. The name came from the C++ pass this file ports, where one method does both -- the docstring now records that so the correspondence is still findable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A Triton kernel describes a single program instance: the tile loop is not in the kernel, it is the launch grid outside it. The trace producer already wants that same split -- togsim_kernel_tile per work-item, enumerated by togsim_kernel (docs/design/togsim_cpp_trace.md sec 9.3) -- so the two models agree and only the enumeration was missing. This teaches the pipeline to accept a kernel in that shape, rather than requiring the loop nest PyTorchSim's codegen emits. build_tog _find_kernel falls back to the module's only func.func; the name comes from the Triton kernel, not the fixed "kernel". _build roots at top-level loops carrying a ROLE attribute, and treats the whole body as one work-item when there are none. Keying on the attribute matters on its own: bank_vectorize also leaves a bare top-level affine.for for the tile's vector work, and rooting at that one made every DMA a sibling the traversal never reached. A DMA's tensor identity comes from the producer's dram_arg attribute when the operand is a view of the argument rather than the argument. Inferring it would mean chasing memref view ops back, and there are nine of them with no ViewLikeOpInterface in the python bindings to ask generically. DMA nodes are recorded in a list so _collect_dma_nodes can seed from them; it only descended from loop nodes, and a DMA outside any loop was dropped before reaching the skeleton. lower_to_emitc WorkItem + _materialize_grid_loop supply the grid, on the trace artifact only: the body is wrapped in a loop per axis, tagged outer_loop, with each program-id argument replaced by its induction variable. It must run before _rewrite_signature, which erases the arguments and first asserts none are still used. Everything after is unchanged -- _parallel_loop_chain finds the tagged loop, the outliner threads the induction variable through iv[], and the loop left behind becomes the dispatch enumeration. Two things about building that nest are easy to get wrong and only show at rank >= 2: a nested loop is created before the enclosing yield (an InsertionPoint on a block appends, and an scf.for body is already terminated), and every bound is created before the first loop, so that a bound made after an outer loop does not end up below it while an inner loop uses it. A parallel loop may be scf.for as well as affine.for; the role is carried by the attribute, not the dialect. _strip_aux keeps the kernel the caller resolved instead of matching on a name. Every change is a fallback: the existing conditions are tested first, so PyTorchSim's own codegen takes exactly the path it did before. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Closes the timing loop: one torch.compile now produces a cycle count. The
launcher emits the trace producer from tnpu's post-vcix IR and hands it to
TOGSim, reusing PyTorchSim's existing trace pipeline unchanged.
timing.py
emit_trace 04-custom.mlir -> build_skeleton -> trace.so + trace_cycles.tsv
run_togsim hand the kernel directory to TOGSimulator.run_standalone
work_item_for derives the program-id argument positions from the signature
layout (pointers, user scalars, then triton-shared's own
gridX,Y,Z / pidX,Y,Z) and the grid from the pinned block sizes
codecache persists meta.json beside the artifacts so the timing step can run
standalone, and TritonNPULauncher.__call__ simulates instead of raising.
kernel_spec._grid becomes grid_of: the timing path needs the same extents to
enumerate work-items, so it is computed in one place.
The test drives a 2-D grid: that it verifies as MLIR, nests one loop per axis,
and dispatches both indices. Checking the module and not only the C++ it becomes
is the point -- the emitc lowering hoists constants to a flat scope, so it hides
a bound that does not dominate its use. Inductor cannot reach this path here (it
uses y/z only when x would overflow), so it gets a test rather than waiting for
a kernel to exercise it.
Measured on Inductor's `x + y` (1024 elements, XBLOCK 128, grid 8): TOGSim
totals 573 cycles, and channel-0 DRAM traffic of 16 reads x 32 B x 16 channels
is 8192 B -- exactly the 8 work-items x 2 loads x 512 B the kernel should move,
so every dispatch in the enumeration really ran. The cycle table is a
placeholder until gem5 sampling lands, and says so on every emit.
Output tensors are NOT written; marshalling them through Spike is the remaining
functional half. The launcher logs that on every call rather than letting an
undefined value pass for a computed one, and the test asserts the timing
artifacts exist instead of comparing values.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The cycle table held a placeholder, so TOGSim modelled the DMA but charged nothing for compute. It is now a measurement. measure_tile_cycles chains three pieces that already existed: build_tog's sample mode annotates the post-vcix IR (inline-asm markers around each compute node, every loop rewritten to one trip, so what runs is one tile), `python -m tnpu.cycle` lowers that to a RISC-V binary in tnpu's own process -- the Gemmini/VCIX lowering and its LLVM live there -- and CycleSimulator runs it under gem5, reading one numCycles per marker pair. build_cycle_table then turns the list into the tsv, keyed by tile_id. Sampling runs before build_skeleton because both consume the same post-vcix IR and build_skeleton rewrites it in place. Failure is not fatal: any step that does not produce a measurement falls back to the placeholder table and says so, since a kernel that simulates with the wrong compute cost is more useful than one that will not simulate -- as long as it announces which it is. Measured on `x + y` (1024 elements, grid 8): the tile samples at 19 cycles and TOGSim's total moves 573 -> 650. The MLIR route reports 251 on the same computation; the remaining gap is double buffering, which tnpu does not emit yet, so nothing overlaps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The grid machinery handles N axes and is tested to, but grid_of only ever
computed one from xnumel, so work_item_for could only ever build a 1-D
WorkItem -- general plumbing behind a caller that never used it.
grid_of now walks every parallel axis, and fixed_config_for pins a block per
axis. Parallel vs reduction is Inductor's own test (a prefix starting with "r"
is looped inside the kernel, not gridded), and the block name is
f"{prefix.upper()}BLOCK", so neither is guessed.
Two orderings meet here and they are not the same: grid_of returns axes
OUTERMOST first (z, y, x -- x is Inductor's contiguous axis), while
triton-shared always appends the program ids as pidX, pidY, pidZ. work_item_for
therefore builds the argument list per axis instead of as a range; zipping the
two blindly would pair the outermost loop with the wrong id.
Block sizes: the outermost axis gets the lane count, because that is the tile
dimension bank_vectorize spreads over the lanes. The rest get 1, which leaves
the tile exactly the verified [lanes] shape and lets the grid cover everything
else. That is correct but pathological -- an inner block of 1 makes each
work-item move a strided column -- so a multi-axis kernel logs a warning saying
it is not a tiling worth measuring. Choosing real tile sizes is the block-size
policy gap in the README.
Verified: axes/grid/parallel_args come out ['x']/(8,)/[pidX],
['y','x']/(2,1024)/[pidY,pidX] and ['z','y','x']/[pidZ,pidY,pidX]. End-to-end on
a real multi-axis kernel is still untested -- Inductor reaches for y/z only when
x would overflow, which the shapes this route handles do not.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
d2a5c9e to
c2e004f
Compare
Only the NUMBER of grid axes has to be compiled in -- how many loops to nest and how many iv[] slots to fill. The trip counts are just values, and the producer ABI already carries them: togsim_kernel(ctx, shape_args, n). Baking them in was what forced a recompile per shape. A WorkItem extent of None now means "read it from shape_args". The bound cannot be wired when the loop is built -- shape_args does not exist until _rewrite_signature adds it -- so the loop takes a placeholder and _bind_runtime_bounds replaces it once the signature is there. The loops stay in the entry function (the outliner moves only their bodies), so the read is in scope where the bound is used. timing.write_shape computes the grid per launch and writes trace_shape.txt. The launch already knows the real extents: Inductor appends the numels after the tensor arguments, so the trailing values are them. Only the PARALLEL numels ride along -- a reduction axis is looped inside the kernel and never passed, so counting it would misalign the mapping. TOGSim reads that sidecar the same way it reads trace_cycles.tsv, from the kernel directory. main.cc passed nullptr for shape_args; absent file still means nullptr, so a producer with its bounds baked in is unaffected. Measured. One trace.so (md5 identical across all three), torch.compile( dynamic=True), a single kernel directory reused: n=1024 grid 8 650 cycles n=2048 grid 16 1316 cycles n=4096 grid 32 2586 cycles The functional path still compiles per shape -- tnpu's spec bakes the tensor extents into the flat memref view -- so this opens the timing half only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 participant
Stacked on #305 (base is
feature/triton-codegen, notdevelop— review that first). Pairs with PSAL-POSTECH/triton-npu#1, which builds the binary the cycle measurement runs.#305 got a Triton kernel compiled to a RISC-V ELF. This gets it simulated: one
torch.compileofx + ynow ends with a TOGSim cycle count.The one structural difference, and why it is small
A Triton kernel describes a single program instance — the tile loop is not in the kernel, it is the launch grid outside it. PyTorchSim's codegen puts the whole loop nest in the kernel, so the two look incompatible.
They are not. The trace producer already splits at exactly that seam (
docs/design/togsim_cpp_trace.md§9.3):togsim_kernel_tileper work-item, enumerated bytogsim_kernel, with core assignment owned bytogsim_dispatch. A Triton kernel arrives in the shape the outliner otherwise has to produce. What was missing was only the enumeration, andlower_to_emitc.WorkItemsupplies it.What changed
build_tog— four fallbacks, each taken only when the existing condition does not hold, so PyTorchSim's own codegen walks the same path as before:_find_kernelaccepts the module's onlyfunc.func(the name comes from the Triton kernel, not the fixedkernel)_buildroots at top-level loops carrying a role attribute, and treats the whole body as one work-item when there are none. Keying on the attribute matters on its own:bank_vectorizealso leaves a bare top-levelaffine.for, and rooting at that one made every DMA a sibling the traversal never reachedmemref.reinterpret_castback to the block argument — triton-shared types pointers as unrankedmemref<*xf32>, so the operand is a flat view of an argument, never the argument, which leftarg_idat −1 and silently broke the address model_collect_dma_nodescan seed from them; it only descended from loop nodes, and a DMA outside any loop was dropped before reaching the skeletonlower_to_emitc—WorkItem+_materialize_grid_loop, applied to the trace artifact only. It runs before_rewrite_signature, which erases the kernel arguments and first asserts none are still used; that ordering is what decides where this can live. Everything after is untouched —_parallel_loop_chainfinds the tagged loop, the outliner threads the induction variable throughiv[], and the loop left behind becomes the dispatch enumeration.triton_backend/timing.py—emit_trace(post-vcix IR →trace.so+trace_cycles.tsv),run_togsim, andmeasure_tile_cycles, which chains three pieces that already existed:build_togsample mode annotates the IR,python -m tnpu.cyclelowers it in tnpu's process, gem5 runs it andCycleSimulatorreads onenumCyclesper marker pair.Measured
x + y, 1024 elements, XBLOCK 128, grid 8:Channel-0 DRAM traffic is 16 reads × 32 B × 16 channels = 8192 B, exactly the 8 work-items × 2 loads × 512 B the kernel should move — so every dispatch in the enumeration really ran.
The MLIR route reports 251 cycles on the same computation. The gap is double buffering: tnpu emits synchronous DMA (
is_async=false, notogsim.wait), so load → compute → store serialize inside every work-item and TOGSim has no overlap to model. That is the next gap, and it is named in the README.What this does not do
Output tensors are not written — marshalling them through Spike is the remaining functional half. The launcher logs that on every call rather than letting an undefined value pass for a computed one, and the test asserts the timing artifacts exist instead of comparing values.
Matmul timing is also still open:
build_togfinds compute nodes by thevcix.ivop name and tnpu emitsllvm.riscv.sf.vc.*intrinsics, so a GEMM currently yields no compute node.Regression
tests/ops/elementwise/test_add.py5/5 andtests/ops/gemm/test_matmul.py11/11 on the MLIR route; tnpu baselinesadderr 0,gemm1.52588e-05. Verified again after rebasing onto the CI commits in #305.